import nbformat
from nbconvert import PythonExporter
import os
import sys


def convert_ipynb_to_py(input_ipynb_path, output_py_path=None):
    """
    将一个 Jupyter Notebook (.ipynb) 文件转换为 Python (.py) 文件。

    Args:
        input_ipynb_path (str): 输入的 .ipynb 文件的路径。
        output_py_path (str, optional): 输出的 .py 文件的路径。
                                    如果为 None，则会在输入文件所在目录下生成同名 .py 文件。
                                    默认为 None。
    """
    if not os.path.exists(input_ipynb_path):
        print(f"错误: 文件不存在: {input_ipynb_path}")
        return

    if output_py_path is None:
        base_name = os.path.splitext(os.path.basename(input_ipynb_path))[0]
        output_py_path = os.path.join(
            os.path.dirname(input_ipynb_path), f"{base_name}.py"
        )

    try:
        with open(input_ipynb_path, "r", encoding="utf-8") as f:
            notebook_content = nbformat.read(f, as_version=4)
    except Exception as e:
        print(f"读取 .ipynb 文件时出错: {e}")
        return

    try:
        exporter = PythonExporter()
        python_code, _ = exporter.from_notebook_node(notebook_content)

        with open(output_py_path, "w", encoding="utf-8") as f:
            f.write(python_code)
        print(f"成功将 {input_ipynb_path} 转换为 {output_py_path}")
    except Exception as e:
        print(f"转换 .ipynb 文件为 .py 文件时出错: {e}")
        return


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("用法: python script.py <input.ipynb> [output.py]")
        print("   <input.ipynb>: 需要转换的 .ipynb 文件路径。")
        print(
            "   [output.py]: 可选参数，输出的 .py 文件路径。如果未指定，则默认生成同名 .py 文件在同一目录下。"
        )
        sys.exit(1)

    input_path = sys.argv[1]
    output_path = None
    if len(sys.argv) > 2:
        output_path = sys.argv[2]

    convert_ipynb_to_py(input_path, output_path)
